// Program to compute average marginal effects after ivprobit.
// Compared to margins this is considerably faster since we use 
// analytic derivatives when calculating the Jacobian matrix (which is
// the derivative of the AME with respect to each model parameter).

// The default is to estimate the marginal effects currently used in stata, which
// are based on the article of Skeeles and Taylor 2015. If the option noendog is
// specified we obtain the marginal effects used before stata 14. 

// Karl-Oskar Lindgren 2017-12-04

 capture program drop mfx_ivprob
 
 program mfx_ivprob, rclass 
  
	preserve 
	
	version 14
	set type double 
	
	syntax varlist (min=1) [, NOENdog] 

quietly {
	
// Make sure that the model current in memory is ivprobit	
	local estcom: word 1 of `e(cmdline)'
	if "`estcom'" != "ivprobit" {
		
		display as error "The latest model was not ivprobit"
		exit
		
	}
	
// Define temporary variables and names	
	tempvar xb ehat jac meff
	tempname rho sigma
	
	predict `xb' if e(sample) == 1 , xb 
  scalar `rho' = tanh(_b[athrho2_1:_cons])
  scalar `sigma' = exp(_b[lnsigma2:_cons])
	
	if ("`noendog'" != "") {
	  scalar `rho' = 0
	}
	
// Save the ivprobit results to memory	
	estimates store _ivprobmod
	
// Generate constant used 	
	capture drop _constant 
	gen _constant = 1
	
// Partion e(b) from ivprobit for later use 	
	matrix b_1 = e(b)
	matrix b_1 = b_1[1, "`e(instd)':"]
	
	local depvar : word 1 of `e(depvar)'
	matrix b_2 = e(b)
	matrix b_2 = b_2[1, "`depvar':"]
		
	local mvars `varlist'
	local nmv : word count `mvars'
	local nk = e(k)

	local vars1 `: colnames b_1'
	local vars1: subinstr local vars1 "_cons" "_cons"
	local vars1 = subinstr("`vars1'", "o.", "", .)
	local vars2 `: colnames b_2'
	local vars2: subinstr local vars2 "_cons" "_cons"
	local vars2 = subinstr("`vars2'", "o.", "", .)
	
// In order to compute the marginal effects a la Skeeles and Taylor we need 
// the "first stage" residuals.	
	qui reg `e(instd)' `vars1' if e(sample) == 1 , nocons
	predict `ehat', res 
	
// Restore the ivprobit results	
	qui: estimates restore _ivprobmod 
	
// Create matrices used for storing	
	matrix Jac = J(`nmv', `nk', 0) 
	matrix mEff = J(`nmv', 1, .)
	matrix rownames mEff = `mvars'
	
	capture drop `jac' 
	capture drop `meff'
	
  local i = 1
  local j = 1
	
// These two variables are used frequently when estimating marginal effects and the Jacobian	
	gen _A = normalden((`xb'+`rho'*`ehat'/`sigma')/(sqrt(1-`rho'^2)))/(sqrt(1-`rho'^2))	 
	gen _B =(`xb'+`rho'*`ehat'/`sigma')/(1-`rho'^2)
	
// Loop over all variables to calculate marginal effects and the Jacobian. 	
	foreach mvar of varlist `mvars' {
	  		
				local checkmvar = strpos("`vars2'", "`mvar'")
				if `checkmvar' == 0 {
				  di as error "The varlist includes `mvar', which does not appear in the 2nd stage of the ivprobit model"
					exit
				}
				
				local eind = "`mvar'" == "`e(instd)'"
		
				if `eind' == 1 {
				  scalar rs = `rho'/`sigma'
				}
				
				else {
				  scalar rs = -1*`rho'/`sigma'*_b["`e(instd)':`mvar'"]
				}
				
				gen `meff' = _A*(_b[`mvar']+rs)
				qui: sum `meff'
				drop `meff' 
				
				matrix mEff[`i', 1] = r(mean)
				
		foreach var2 of varlist `vars2' {
	
		    local ind  = "`mvar'" == "`var2'"
			
		    gen `jac' = _A*(`ind'-`var2'*(_b[`mvar']+rs)*_B)
				qui: sum `jac'

        matrix Jac[`i', `j'] = r(mean)
			  drop `jac'
        local j = `j'+1
	}	
	  foreach var1 of varlist `vars1' {
		    
				local eind = "`mvar'" == "`e(instd)'"
		    local ind1  = "`mvar'" == "`var1'"
				
				if `eind' == 1 {
				  scalar rs = `rho'/`sigma'
				}
				
				else {
				  scalar rs = -1*`rho'/`sigma'*_b["`e(instd)':`mvar'"]
				}
				
		    gen `jac' = _A*(-`ind1'*(`rho'/`sigma')+`var1'*(`rho'/`sigma')*(_b[`mvar']+rs)*_B)
				qui: sum `jac'

        matrix Jac[`i', `j'] = r(mean)
			  drop `jac'
        local j = `j'+1
		}
	  local j = 1
    local i = `i'+1
	}
	
	  matrix vCov = e(V)
	  matrix mV = vecdiag(Jac*vCov*Jac')'
	  
	  matrix OutMat = J(`nmv', 4, .)
	  matrix rownames OutMat = `mvars'
	  matrix colnames OutMat = "dy/dx" "Std. Err." "z" "P>|z|"   
		
	foreach num of numlist 1/`nmv' {
	  
		matrix OutMat [`num', 1] = mEff[`num', 1]
		matrix OutMat [`num', 2] = sqrt(mV[`num', 1])
		matrix OutMat [`num', 3] = mEff[`num', 1]/sqrt(mV[`num', 1])
		matrix OutMat [`num', 4] = (1-normal(abs(mEff[`num', 1]/sqrt(mV[`num', 1]))))*2
	}
	
	matrix mfx_b = OutMat[1..`nmv', 1]
	matrix mfx_se = OutMat[1..`nmv', 2] 
	return matrix mfx_b = mfx_b 
	return matrix mfx_se = mfx_se 
}
	di ""
	di as text "Average marginal effects           Number of obs = " as result `e(N)'
	di as text "Delta method Std. Err."
    di ""
    matlist OutMat, border(rows) noheader 
	
end
